from sklearn import datasets
from losses import CrossEntropy
import numpy as np
import random
import matplotlib.pyplot as plt

from layers import LinearLayer, SigmoidLayer
from sklearn import model_selection


iris = datasets.load_iris()

features = np.array(iris['data'])
targets = np.array(iris['target'])

num_samples = len(features)

labels = np.zeros((num_samples,3))
for i in range(num_samples):
    labels[i, targets[i]] = 1


def data_iter(batch_size, features, labels):
    num_examples = len(features)
    indices = list(range(num_examples))
    # 这些样本是随机读取的，没有特定的顺序
    random.shuffle(indices)
    for i in range(0, num_examples, batch_size):
        batch_indices = indices[i: min(i + batch_size, num_examples)]
        yield features[batch_indices], labels[batch_indices]


# ================ 网络 ==================
net = [
    LinearLayer(4, 6),
    SigmoidLayer(),
    LinearLayer(6, 3),
    SigmoidLayer()
]

def forward(x):
    """正向传播"""
    out = x
    for layer in net:
        out = layer._forward(out)

    return out


def backward(grad, lr):
    """反向传播"""
    grad_prev = grad
    for layer in reversed(net):
        grad_prev = layer._backward(grad_prev, lr=lr)

def accuracy(y_hat, y):
    right_count=0
    for i in range(len(y)):
        if np.argmax(y_hat[i]) == np.argmax(y[i]):
            right_count+=1

    return right_count/ len(y)



# =============== 训练 =================
lr = 0.05
loss = CrossEntropy

epochs = 5000
train_losses = []
test_acc = []
for ep in range(epochs):
    X_train, X_test, Y_train, Y_test = model_selection.train_test_split(features, labels)

    for x, y in data_iter(28, X_train, Y_train):
        y_hat = forward(x)

        e = loss.fn(y_hat, y)
        grad_e_y = loss.grad(y_hat, y)

        backward(grad_e_y, lr)

    train_l = loss.fn(forward(X_train), Y_train).mean()
    train_losses.append(train_l)
    
    y_hat_test = forward(X_test)
    acc = accuracy(y_hat_test, Y_test)
    test_acc.append(acc)
    print("epoch:{}, loss:{:.4f}, acc:{}\n".format(ep, train_l, acc))

    if train_l < 0.4: lr = 0.1
    if train_l < 0.15 and acc > 0.97: break    


total_acc = accuracy(forward(features), labels)
print("total accuracy:", total_acc)

plt.plot(train_losses, label = "loss")
plt.plot(test_acc, label = "accuracy")
plt.legend()
plt.show()







